import os, sys
import json
import tqdm
path = "."
json_list = []
idx = 0
for jfile in os.listdir(path):
    #if idx > 0:
    #    break
    #if not jfile.endswith("N-smooth-numbers.json"):
    #    continue
    json_list.append(jfile)
    idx += 1

def get_code_snippet(code_snippet):
    rust_cs = list()
    c_cs = list()
    for idx, cs in code_snippet.items():
        #cs = code_snippet[idx]
        if "rust" in cs:
            rust_cs.append(cs["rust"])
        if "c" in cs:
            c_cs.append(cs["c"])
    return rust_cs, c_cs

def same_snippet(rust_k, rust_n, c_k, c_n):
    global_kind = ("const", "static")
    if rust_n != c_n:
        return False
    if rust_k == c_k or (rust_k == "fn" and c_k == "function"):
        return True
    elif rust_k in global_kind and c_k in global_kind:
        return True
    else:
        return False

def get_code_pair(rust_cs, c_cs):
    main_lst = []
    tmp_lst = []
    for c_d in c_cs:
        c_c = c_d["code"]
        c_k = c_d["kind"]
        c_n = c_d["name"]
        c_f = c_d["file"]
        if c_k == "function" and c_n == "main":
            main_lst.append(c_d)
            tmp_lst.append([])
        else:
            tmp_lst.append(c_d)
    
    main_dict = main_lst[0]
    code = ""
    for c_d in main_lst:
        code += c_d["code"]
    main_dict["code"] = code
    main_dict["name"] = "main"
    flag = False
    c_cs = []
    for idx in range(len(tmp_lst)):
        if not tmp_lst[idx]:
            if flag == False:
                tmp_lst[idx] = main_dict
                flag = True
                c_cs.append(tmp_lst[idx])
        else:
            c_cs.append(tmp_lst[idx])
    #print(c_cs)
    rust_idx = 0
    rust_s = 0
    c_s = 0
    same_d = []
    idx = 0
    resort = []
    while rust_idx < len(rust_cs):
        rust_d = rust_cs[rust_idx]
        rust_c = rust_d["code"]
        rust_k = rust_d["kind"]
        rust_n = rust_d["name"]
        rust_f = rust_d["file"]
        flag = False
        c_idx = 0
        while c_idx < len(c_cs):
            c_d = c_cs[c_idx]
            c_c = c_d["code"]
            c_k = c_d["kind"]
            c_n = c_d["name"]
            c_f = c_d["file"]
            if same_snippet(rust_k, rust_n, c_k, c_n):
                flag = True
                break
            c_idx += 1 
        if flag:
            resort.append(list())
            rust_flag = False
            c_flag = False
            while rust_s < rust_idx:
                rust_flag = True
                resort[idx].append(("rust", rust_cs[rust_s]))
                rust_s += 1
            while c_s < c_idx:
                c_flag = True
                resort[idx].append(("c", c_cs[c_s]))
                c_s += 1
            if rust_flag or c_flag:
                idx += 1
            resort.append(list())
            resort[idx].append(("rust", rust_d))
            resort[idx].append(("c", c_d))
            idx += 1
            #print(rust_s, c_s, rust_idx, c_idx)
            #rust_s += 1 if rust_flag else 1
            #c_s += 1 if c_flag else 1
            rust_s += 1
            c_s += 1
        rust_idx += 1
        #print(rust_s, c_s, rust_idx, c_idx)
    merge_dict = dict()
    num = 1
    i = 0
    while i < idx:
        item = resort[i]
        merge_dict[str(num)] = dict()
        rust_c = ""
        c_c = ""
        for d in item:
            if d[0] == "rust":
                rust_c += d[1]["code"]
            elif d[0] == "c":
                c_c += d[1]["code"]
            else:
                print(f"Unexpcted item:{d}")
        if rust_c != "" or c_c != "":
            merge_dict[str(num)]["rust"] = rust_c
            merge_dict[str(num)]["c"] = c_c
            num += 1
        i += 1
    return merge_dict

json_list = sorted(json_list)
for jfile in tqdm.tqdm(json_list):
  if jfile.endswith(".json") and not jfile.endswith("_new.json"):
    print(f"Parsing {jfile}")
    with open(jfile, "r", encoding="utf-8") as j:
        content = j.read()
        code_snippet = json.loads(content)
        rust_cs, c_cs = get_code_snippet(code_snippet)
        pair = get_code_pair(rust_cs, c_cs)
        new_jfile = jfile.replace(".json", "_new.json")
        print(f"new json:{new_jfile}")
        with open(new_jfile, "w", encoding="utf-8") as newj:
            json.dump(pair, newj, ensure_ascii=False, indent=4)
